import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LassoCV,Ridge
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# 设置字体以支持 Unicode 字符
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 读取数据
data = pd.read_excel("data/result.xlsx")
data = data.fillna(0)

y = data['销量(千克)']
X = data[['销售单价(元/千克)', '销售金额']]  # 确保 X 是一个 DataFrame

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 使用 LassoCV 自动选择最佳 alpha
lasso_cv = LassoCV(alphas=np.logspace(-4, 0, 50), cv=5)
lasso_cv.fit(X_train, y_train)

# 打印最终决定的系数
print(f'最终决定的系数: {lasso_cv.coef_}')
print(f'最佳 alpha: {lasso_cv.alpha_}')

# 绘制 Lasso 回归系数选择图
alphas = np.logspace(-4, 0, 50)
coefs = []

for alpha in alphas:
    lasso = LassoCV(alphas=[alpha], cv=5)
    lasso.fit(X_train, y_train)
    coefs.append(lasso.coef_)

plt.figure(figsize=(10, 6))
plt.plot(alphas, coefs)
plt.xscale('log')
plt.xlabel('Alpha')
plt.ylabel('Coefficients')
plt.title('Lasso 回归系数选择图')
plt.axis('tight')
plt.grid()
plt.show()

# 获取 Lasso 选择的非零系数特征
selected_features = X_train.columns[lasso_cv.coef_ != 0]

# 使用 Ridge 回归对选择的特征进行拟合
ridge = Ridge(alpha=1.0)  # 这里的 alpha 可以根据需要调整
ridge.fit(X_train[selected_features], y_train)

# 预测
y_pred_train = ridge.predict(X_train[selected_features])
y_pred_test = ridge.predict(X_test[selected_features])

# 计算均方误差
mse_train = mean_squared_error(y_train, y_pred_train)
mse_test = mean_squared_error(y_test, y_pred_test)

print(f'训练集 MSE: {mse_train}')
print(f'测试集 MSE: {mse_test}')

# 绘制预测效果图
plt.figure(figsize=(10, 6))
plt.scatter(y_test, y_pred_test, color='blue', label='预测值')
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=2, label='理想值')
plt.xlabel('实际值')
plt.ylabel('预测值')
plt.title('岭回归预测效果图')
plt.legend()
plt.show()

# 绘制3D预测效果图
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')

# 绘制实际值
ax.scatter(X_test['销售单价(元/千克)'], X_test['销售金额'], y_test, color='blue', label='实际值')

# 绘制预测值
ax.scatter(X_test['销售单价(元/千克)'], X_test['销售金额'], y_pred_test, color='red', label='预测值')

# 绘制拟合平面
x_surf, y_surf = np.meshgrid(np.linspace(X_test['销售单价(元/千克)'].min(), X_test['销售单价(元/千克)'].max(), 100),
                             np.linspace(X_test['销售金额'].min(), X_test['销售金额'].max(), 100))
z_surf = ridge.predict(np.c_[x_surf.ravel(), y_surf.ravel()]).reshape(x_surf.shape)

ax.plot_surface(x_surf, y_surf, z_surf, color='yellow', alpha=0.5, label='拟合平面')

ax.set_xlabel('销售单价(元/千克)')
ax.set_ylabel('销售金额')
ax.set_zlabel('销量(千克)')
ax.set_title('岭回归预测效果图')
ax.legend()
plt.show()